import os
from pydub import AudioSegment

def split_audio(input_audio_file, output_prefix,speak_name, start_sec, end_sec):
    # 加载音频文件
    audio = AudioSegment.from_file(input_audio_file)

    # 将开始秒数和结束秒数转换为毫秒
    start_ms = start_sec
    end_ms = end_sec

    # 切分音频
    split_audio = audio[start_ms*1000:end_ms*1000]

    # 保存切分后的音频文件
    output_file=os.path.join(output_prefix, f"{speak_name}_{start_sec}s_to_{end_sec}s.wav")

    split_audio.export(output_file, format="wav")
    return output_file

def batch_split_audio(input_folder, output_folder, start_sec, end_sec):
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历输入文件夹中的所有文件
    for file_name in os.listdir(input_folder):
        if file_name.endswith('.wav'):  # 只处理 WAV 文件，可以根据需要修改
            # 构建输入文件的完整路径
            input_file_path = os.path.join(input_folder, file_name)

            # 构建输出文件的完整路径
            output_file_prefix = os.path.splitext(file_name)[0]
            output_file_folder = os.path.join(output_folder, output_file_prefix)
            if not os.path.exists(output_file_folder):
                os.makedirs(output_file_folder)

            # 执行切分操作
            split_audio(input_file_path, os.path.join(output_file_folder, output_file_prefix), start_sec, end_sec)

def merge_short_segments(input_folder, output_folder, target_duration_ms=5000):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    # 初始化新的音频片段列表
    new_segments = []
    short_length = 0
    # 遍历输入文件夹中的所有文件
    for filename in os.listdir(input_folder):
        # 检查文件是否为音频文件
        if filename.endswith(".wav") or filename.endswith(".WAV"):
            input_audio_file = os.path.join(input_folder, filename)
            output_audio_file = os.path.join(output_folder, filename)

            # 加载音频文件
            audio = AudioSegment.from_file(input_audio_file)

            if len(audio) < target_duration_ms:
                # 如果音频文件长度小于目标长度，则添加到新的音频片段列表中
                new_segments.append(audio)
                short_length += len(audio)
                if short_length >= target_duration_ms:
                    # 将新的音频片段合并成一个音频文件
                    merged_audio = AudioSegment.empty();
                    for j, audio_segment in enumerate(new_segments):
                        print(f"合并音频段{j + 1}")
                        merged_audio += audio_segment

                    merged_audio.set_frame_rate(44100)
                    print(f"音频文件{filename}的长度小于{target_duration_ms}ms，合并后的长度为{len(merged_audio)}ms")
                    merged_audio.export(output_audio_file, format="wav")

                    new_segments = []
                    short_length = 0
            else:
                print(f"音频文件{filename}的长度大于{target_duration_ms}ms，直接复制")
                audio.export(output_audio_file, format="wav")

def split_audio_into_segments(input_audio_file, output_folder,speak_name, segment_length_ms):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    # 加载音频文件
    audio = AudioSegment.from_file(input_audio_file)

    # 计算每个片段的起始和结束时间
    start_time = 0
    end_time = segment_length_ms

    # 循环切分音频
    segment_index = 1
    while start_time < len(audio):
        # 切分音频片段
        segment = audio[start_time:end_time]

        # 保存切分后的音频片段
        output_file = f"{output_folder}/{speak_name}_segment_{segment_index}.wav"
        segment.export(output_file, format="wav")

        # 更新起始和结束时间
        start_time = end_time
        end_time += segment_length_ms

        # 更新片段索引
        segment_index += 1


merge_short_segments("../data/lanying/vad/huanzou_7db", "../data/lanying/vad/huanzou_7db_merge",5000)
